import org.apache.log4j.{Level, Logger}
import org.apache.spark.sql.{DataFrame, Dataset, Row, SparkSession}
import scala.math.{pow, sqrt}

/**
 * 4、使用鸢尾花数据集实现KNN算法
 *    文件A（data/Iris_A.csv)：141条数据，分类已知
 *    文件B（data/Iris_B.csv)：9条数据，分类假设为未知
 *
 * 读文件A、B，形成数据集X、Y
 * 求数据集Y中每个点，到数据集X中每个点的距离，得到数据集D
 * 找到数据集D中距离最小的K个点
 * 求K个点的分类情况
 * 距离使用欧几里得距离
 */

object Subject4 {
  def main(args: Array[String]): Unit = {

    Logger.getLogger("org").setLevel(Level.WARN)
    val spark = SparkSession
      .builder()
      .appName(this.getClass.getCanonicalName)
      .master("local[*]")
      .getOrCreate()
    import spark.implicits._
    import spark._

    //设置K值
    val K = 9

    //读取已知分类的样本数据
    val dfA: DataFrame = spark.read
      .option("header", "true")
      .option("delimiter", ",")
      .option("inferschema", "true")
      .csv("data/Iris_A.csv")

    //读取未知分类的测试数据
    val rowsB: Array[Row] = spark.read
      .option("header", "true")
      .option("delimiter", ",")
      .option("inferschema", "true")
      .csv("data/Iris_B.csv")
      .collect()

    //计算两点的欧氏距离
    def getDistance(slA: Double, swA: Double, plA: Double, pwA: Double,
                    slB: Double, swB: Double, plB: Double, pwB: Double) : Double = {
      sqrt(pow(slA - slB, 2) + pow(swA - swB, 2) + pow(plA - plB, 2) + pow(pwA - pwB, 2))
    }

    //求测试数据中的每个点到样本数据中每个点的距离，取距离最小的K个点，预测样本数据的分类
    rowsB.foreach{ row =>
      //(测试数据id, 两点间距, 对应样本数据的分类)
      val ds: Dataset[(Int, Double, String)] = dfA.map { line =>
        (row.getInt(0),
          getDistance(line.getDouble(1), line.getDouble(2), line.getDouble(3), line.getDouble(4),
            row.getDouble(1), row.getDouble(2), row.getDouble(3), row.getDouble(4)),
          line.getString(5))
      }
      //根据两点间距排序，取间距最小的K个点，预测样本数据分类
      ds.sort('_2).take(K).foreach(println)
      ds.sort('_2).take(K).groupBy(_._3).mapValues(_.length).foreach(println)
      println("****************************************")
    }

  }
}
